#!/usr/bin/env python3
"""MCP Server for Hydrology Forecast Analysis."""

import asyncio
import json
from typing import Any, Sequence
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent

from .forecast_client import ForecastAPIClient
from .forecast_analyzer import ForecastAnalyzer


# Initialize server
app = Server("hydrology-forecast-server")

# Global client instance
client = ForecastAPIClient()


@app.list_tools()
async def list_tools() -> list[Tool]:
    """List available tools."""
    return [
        Tool(
            name="fetch_forecast_session",
            description="""获取水文预报会话的基本信息。
            
            输入参数：
            - session_id: 预报会话ID（必填）
            - use_cache: 是否使用缓存数据（可选，默认true）
            
            返回：站点信息、模型类型、假拟预报数量等基本信息。
            """,
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {
                        "type": "string",
                        "description": "预报会话ID"
                    },
                    "use_cache": {
                        "type": "boolean",
                        "description": "是否使用缓存数据",
                        "default": True
                    }
                },
                "required": ["session_id"]
            }
        ),
        Tool(
            name="analyze_flood_metrics",
            description="""分析水文预报的洪水指标（洪量、洪峰、峰现时间）。
            
            输入参数：
            - session_id: 预报会话ID（必填）
            - forecast_type: 预报类型，可选值：
              * "real" - 仅分析真实预报
              * "hypothetical" - 分析所有假拟预报
              * "all" - 分析所有预报（默认）
            - hypothetical_index: 指定分析某个假拟预报的索引（可选）
            
            返回：洪量（m³）、洪峰流量（m³/s）、峰现时间等指标。
            """,
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {
                        "type": "string",
                        "description": "预报会话ID"
                    },
                    "forecast_type": {
                        "type": "string",
                        "enum": ["real", "hypothetical", "all"],
                        "description": "要分析的预报类型",
                        "default": "all"
                    },
                    "hypothetical_index": {
                        "type": "integer",
                        "description": "假拟预报索引（仅当forecast_type为hypothetical时使用）"
                    }
                },
                "required": ["session_id"]
            }
        ),
        Tool(
            name="compare_forecasts",
            description="""对比真实预报与假拟预报的差异。
            
            输入参数：
            - session_id: 预报会话ID（必填）
            
            返回：详细的对比结果，包括：
            - 洪量差异（绝对值和百分比）
            - 洪峰流量差异（绝对值和百分比）
            - 峰现时间差异
            - 统计摘要
            """,
            inputSchema={
                "type": "object",
                "properties": {
                    "session_id": {
                        "type": "string",
                        "description": "预报会话ID"
                    }
                },
                "required": ["session_id"]
            }
        )
    ]


@app.call_tool()
async def call_tool(name: str, arguments: Any) -> Sequence[TextContent]:
    """Handle tool calls."""
    try:
        if name == "fetch_forecast_session":
            return await handle_fetch_session(arguments)
        elif name == "analyze_flood_metrics":
            return await handle_analyze_metrics(arguments)
        elif name == "compare_forecasts":
            return await handle_compare_forecasts(arguments)
        else:
            return [TextContent(type="text", text=f"Unknown tool: {name}")]
    except Exception as e:
        return [TextContent(type="text", text=f"Error: {str(e)}")]


async def handle_fetch_session(args: dict) -> Sequence[TextContent]:
    """Handle fetch_forecast_session tool."""
    session_id = args.get("session_id")
    use_cache = args.get("use_cache", True)
    
    # Fetch data
    data = await client.fetch_forecast_data(session_id, use_cache)
    
    # Analyze
    analyzer = ForecastAnalyzer(data)
    session_info = analyzer.get_session_info()
    model_params = analyzer.get_model_parameters()
    
    result = {
        "session_info": session_info,
        "model_parameters": model_params
    }
    
    return [TextContent(
        type="text",
        text=json.dumps(result, ensure_ascii=False, indent=2)
    )]


async def handle_analyze_metrics(args: dict) -> Sequence[TextContent]:
    """Handle analyze_flood_metrics tool."""
    session_id = args.get("session_id")
    forecast_type = args.get("forecast_type", "all")
    hypothetical_index = args.get("hypothetical_index")
    
    # Fetch data
    data = await client.fetch_forecast_data(session_id)
    analyzer = ForecastAnalyzer(data)
    
    result = {}
    
    if forecast_type in ["real", "all"]:
        result["real_forecast"] = analyzer.analyze_real_forecast()
    
    if forecast_type in ["hypothetical", "all"]:
        if hypothetical_index is not None:
            result["hypothetical_forecast"] = analyzer.analyze_hypothetical_forecast(hypothetical_index)
        else:
            result["hypothetical_forecasts"] = analyzer.analyze_all_hypothetical()
    
    return [TextContent(
        type="text",
        text=json.dumps(result, ensure_ascii=False, indent=2)
    )]


async def handle_compare_forecasts(args: dict) -> Sequence[TextContent]:
    """Handle compare_forecasts tool."""
    session_id = args.get("session_id")
    
    # Fetch data
    data = await client.fetch_forecast_data(session_id)
    analyzer = ForecastAnalyzer(data)
    
    # Compare
    comparison = analyzer.compare_forecasts()
    
    return [TextContent(
        type="text",
        text=json.dumps(comparison, ensure_ascii=False, indent=2)
    )]


async def main():
    """Run the MCP server."""
    async with stdio_server() as (read_stream, write_stream):
        await app.run(
            read_stream,
            write_stream,
            app.create_initialization_options()
        )


if __name__ == "__main__":
    asyncio.run(main())

